import cv2
import os
import shutil
import numpy as np
import sys

# 从命令行参数读取源目录
if len(sys.argv) < 2:
    print("使用方法: python script.py <source_dir>")
    sys.exit(1)

source_dir = sys.argv[1]  # 输入文件夹
out_dir = os.path.join(source_dir, "out")  # 输出文件夹
to_dir = os.path.join(out_dir, '已遮脸')  # 检测到人脸的输出文件夹
noto_dir = os.path.join(out_dir, '未遮脸')  # 未检测到人脸的输出文件夹
# 获取当前脚本的路径
script_dir = os.path.dirname(os.path.abspath(__file__))

# 加载 OpenCV 的 DNN 模型 (Caffe)
# 定义 prototxt 文件和 caffe 模型文件的相对路径
prototxt_relative_path = "deploy.prototxt"
caffemodel_relative_path = "res10_300x300_ssd_iter_140000.caffemodel"
# 构建绝对路径
prototxt_path = os.path.join(script_dir, prototxt_relative_path)
caffemodel_path = os.path.join(script_dir, caffemodel_relative_path)
face_mask_pic = os.path.join(script_dir, 'face_mask.png')
# 确保目标文件夹存在
for directory in [to_dir, noto_dir]:
    if not os.path.exists(directory):
        os.makedirs(directory)


 
 
# 构建绝对路径
prototxt_path = os.path.join(script_dir, prototxt_relative_path)
caffemodel_path = os.path.join(script_dir, caffemodel_relative_path)
net = cv2.dnn.readNetFromCaffe(prototxt_path, caffemodel_path)

# 加载遮挡用的 PNG 图片（带透明背景的遮脸图片）
overlay = cv2.imread(face_mask_pic , cv2.IMREAD_UNCHANGED)

# 读取支持中文路径的图片
def imread_with_chinese_path(path):
    img_data = np.fromfile(path, dtype=np.uint8)
    img = cv2.imdecode(img_data, cv2.IMREAD_UNCHANGED)
    return img

# 遍历源目录中的所有文件
for filename in os.listdir(source_dir):
    # 只处理图片文件
    if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
        img_path = os.path.join(source_dir, filename)

        # 读取图像
        img = imread_with_chinese_path(img_path)

        # 检查图像是否成功加载
        if img is None:
            print(f"无法加载图像: {img_path}")
            continue

        # 获取图像的高度和宽度
        h, w = img.shape[:2]

        # 将图像转换为适合 DNN 输入的 blob 格式
        blob = cv2.dnn.blobFromImage(cv2.resize(img, (300, 300)), 1.0, (300, 300), (104.0, 177.0, 123.0))

        # 前向传播，进行人脸检测
        net.setInput(blob)
        detections = net.forward()

        # 标志变量，判断是否检测到人脸
        face_detected = False

        # 遍历检测结果
        for i in range(detections.shape[2]):
            confidence = detections[0, 0, i, 2]

            if confidence > 0.5:  # 设置检测置信度阈值
                face_detected = True  # 设置标志变量
                box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
                (x, y, x1, y1) = box.astype("int")

                # 增加边距
                margin = 50
                x = max(0, x - margin)
                y = max(0, y - margin)
                x1 = min(w, x1 + margin)
                y1 = min(h, y1 + margin)

                # 调整遮挡图片的大小以匹配脸部大小
                face_w = x1 - x
                face_h = y1 - y
                resized_overlay = cv2.resize(overlay, (face_w, face_h))

                # 分离遮挡图片的 BGR 与 Alpha 通道
                overlay_rgb = resized_overlay[:, :, :3]
                overlay_alpha = resized_overlay[:, :, 3] / 255.0  # Alpha 通道归一化

                # 获取脸部区域的 ROI
                roi = img[y:y1, x:x1]

                # 合成遮挡图片和原图的 ROI 区域
                for c in range(0, 3):
                    roi[:, :, c] = (overlay_alpha * overlay_rgb[:, :, c] + (1 - overlay_alpha) * roi[:, :, c])

                # 将处理后的 ROI 回填到原图
                img[y:y1, x:x1] = roi

        # 根据是否检测到人脸保存图像
        if face_detected:
            output_path = os.path.join(to_dir, filename)
            cv2.imencode('.jpg', img)[1].tofile(output_path)
            print(f"检测到人脸，已保存到: {output_path}")
        else:
            noface_path = os.path.join(noto_dir, filename)
            shutil.copy(img_path, noface_path)
            print(f"未检测到人脸，已复制到: {noface_path}")
success_file_path = os.path.join(source_dir, 'success.txt')

with open(success_file_path, 'w', encoding='utf-8') as f:
    f.write("所有图像已成功处理，并保存到目标目录。")

print(f"处理完成！所有图像已保存到目标目录。\nsuccess.txt 文件已生成于: {success_file_path}")
 
